﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Runtime.InteropServices;
using System.Threading;
using JudgeKernel.Judge;

namespace JudgeKernel
{
	internal class JudgeProcessor
	{

		#region 私有数据区域

		IntPtr MainProcess = IntPtr.Zero;
		IntPtr MainThread = IntPtr.Zero;
		IntPtr WriteStream = IntPtr.Zero;

		IntPtr ReadStream = IntPtr.Zero;
		IntPtr job = IntPtr.Zero;

		IntPtr ErrorStream = IntPtr.Zero;
		IntPtr IoPort = IntPtr.Zero;
		#endregion

		#region 属性区域

		/// <summary>
		/// 程序路径
		/// </summary> 
		public string ExePath { get; set; }

		/// <summary>
		/// 程序参数
		/// </summary>
		public string ExeParam { get; set; }

		/// <summary>
		/// 用户名密码
		/// </summary>
		public KeyValuePair<string, string>? Certify { get; set; }

		int _MemoryLimit = -1;
		/// <summary>
		/// 内存限制
		/// </summary>
		public int MemoryLimit { get { return _MemoryLimit; } set { _MemoryLimit = value; } }

		int _TimeLimit = -1;
		/// <summary>
		/// 时间限制
		/// </summary>
		public int TimeLimit { get { return _TimeLimit; } set { _TimeLimit = value; } }

		int _OutputLimit = -1;
		/// <summary>
		/// 输出缓冲区大小限制
		/// </summary>
		public int OutputLimit { get { return _OutputLimit - 1024; } set { _OutputLimit = value + 1024; } }

		/// <summary>
		/// 测试数据
		/// </summary>
		public String InputData { get; set; }

		List<byte> _OutputData = new List<byte>();
		/// <summary>
		/// 最后的结果(只读的)
		/// </summary>
		public String OutputData { get { return Encoding.Default.GetString(_OutputData.ToArray()); } }

		uint _Error;
		/// <summary>
		/// 运行时错误
		/// </summary>
		public uint RuntimeErrorState { get { return _Error; } }

		/// <summary>
		/// 程序运行时间(只读的)
		/// </summary>
		public uint TimeUsed
		{
			get
			{
				WIN32.FILETIME c, e, k, u;
				WIN32.GetProcessTimes(MainProcess, out c, out e, out k, out u);
				WIN32.SYSTEMTIME sk, su;
				WIN32.FileTimeToSystemTime(ref k, out sk);
				WIN32.FileTimeToSystemTime(ref u, out su);
				int result = sk.wMilliseconds + sk.wSecond * 1000 + su.wSecond * 1000 + su.wMilliseconds - 15;
				return (uint)(result <= 15 ? 0 : result);
			}
		}

		JudeResult _Status = JudeResult.OK;
		Object __Status = new Object();
		/// <summary>
		/// 运行结果(只读的)
		/// </summary>
		public JudeResult Status { get { return _Status; } }

		/// <summary>
		/// 内存峰值(只读的)
		/// </summary>
		public uint MemoryUsed
		{
			get
			{
				WIN32.PROCESS_MEMORY_COUNTERS pmifo = new WIN32.PROCESS_MEMORY_COUNTERS { cb = (uint)Marshal.SizeOf(typeof(WIN32.PROCESS_MEMORY_COUNTERS)) };
				WIN32.GetProcessMemoryInfo(MainProcess, out pmifo, (uint)Marshal.SizeOf(typeof(WIN32.PROCESS_MEMORY_COUNTERS)));
				return pmifo.PeakPagefileUsage;
			}
		}

		#endregion

		public JudgeProcessor()
		{
		}

		/// <summary>
		/// 启动任务,并开始执行
		/// </summary>
		public bool Judge()
		{
			if (_OutputLimit == -1 || ExePath == null || InputData == null || _MemoryLimit == -1 || _TimeLimit == -1) {
				return false;
			}

			WIN32.STARTUPINFOW stinfo = new WIN32.STARTUPINFOW { cb = (uint)Marshal.SizeOf(typeof(WIN32.STARTUPINFOW)) };
			WIN32.PROCESS_INFORMATION pinfo;
			WIN32.SECURITY_ATTRIBUTES sa;
			sa.bInheritHandle = true;
			sa.lpSecurityDescriptor = IntPtr.Zero;
			sa.nLength = (uint)Marshal.SizeOf(typeof(WIN32.SECURITY_ATTRIBUTES));

			#region 创建重定向管道

			IntPtr hReadStdin, hWriteStdin;
			IntPtr hReadStdout, hWriteStdout;

			WIN32.CreatePipe(out hReadStdin, out hWriteStdin, ref sa, (uint)(InputData.Length + 1));
			WIN32.CreatePipe(out hReadStdout, out hWriteStdout, ref sa, 0);
			WIN32.DuplicateHandle(WIN32.GetCurrentProcess(), hWriteStdin, WIN32.GetCurrentProcess(), out WriteStream,
								  0, false, WIN32.DUPLICATE_SAME_ACCESS);
			WIN32.DuplicateHandle(WIN32.GetCurrentProcess(), hReadStdout, WIN32.GetCurrentProcess(), out ReadStream,
								  0, false, WIN32.DUPLICATE_SAME_ACCESS);


			WIN32.CloseHandle(hWriteStdin);
			WIN32.CloseHandle(hReadStdout);
			uint tmp;
			var inputBytes = Encoding.Default.GetBytes(InputData);
			WIN32.WriteFile(WriteStream, inputBytes, (uint)inputBytes.Length, out tmp, IntPtr.Zero);

			stinfo.dwFlags = WIN32.STARTF_USESTDHANDLES;
			stinfo.hStdInput = hReadStdin;
			stinfo.hStdOutput = hWriteStdout;

			#endregion

			WIN32.CreateProcess(null, ExePath, IntPtr.Zero, IntPtr.Zero, true,
				/*WIN32.CREATE_NEW_CONSOLE & WIN32.CREATE_BREAKAWAY_FROM_JOB &*/ WIN32.CREATE_SUSPENDED,
								null, null, ref stinfo, out pinfo);

			MainProcess = pinfo.hProcess;//
			MainThread = pinfo.hThread;

			if (MemoryUsed > _MemoryLimit) {
				_Status = JudeResult.MemoryLimitExceeded;
				WIN32.TerminateProcess(MainProcess, 0);
				WIN32.WaitForSingleObject(MainProcess, WIN32.INFINITE);
				Thread.Sleep(50);
				return true;
			}


			#region 创建sandbox

			job = WIN32.CreateJobObject(IntPtr.Zero, "joboject");

			//内存设置,并设置取消弹出异常对话框
			WIN32.JOBOBJECT_EXTENDED_LIMIT_INFORMATION basc = new WIN32.JOBOBJECT_EXTENDED_LIMIT_INFORMATION { ProcessMemoryLimit = (uint)_MemoryLimit };
			basc.BasicLimitInformation.LimitFlags = WIN32.JOB_OBJECT_LIMIT_PROCESS_MEMORY | WIN32.JOB_OBJECT_LIMIT_DIE_ON_UNHANDLED_EXCEPTION;
			WIN32.SetInformationJobObject(job, JOBOBJECTINFOCLASS.JobObjectExtendedLimitInformation, ref basc,
										  (uint)Marshal.SizeOf(typeof(WIN32.JOBOBJECT_EXTENDED_LIMIT_INFORMATION)));


			WIN32.JOBOBJECT_ASSOCIATE_COMPLETION_PORT port;
			IoPort = WIN32.CreateIoCompletionPort(new IntPtr(-1), new IntPtr(0), 0, 0);
			port.CompletionPort = IoPort;
			port.CompletionKey = new IntPtr(1);
			WIN32.SetInformationJobObject(job, JOBOBJECTINFOCLASS.JobObjectAssociateCompletionPortInformation,
										  ref port, (uint)Marshal.SizeOf(typeof(WIN32.JOBOBJECT_ASSOCIATE_COMPLETION_PORT)));


			WIN32.AssignProcessToJobObject(job, MainProcess);
			#endregion


			Thread run1 = new Thread(new ThreadStart(Thread1));
			Thread run2 = new Thread(new ThreadStart(Thread2));
			Thread run3 = new Thread(new ThreadStart(Thread3));
			run1.Start();
			run2.Start();
			run3.Start();
			WIN32.ResumeThread(MainThread);

            WIN32.CloseHandle(WriteStream);
			WIN32.WaitForSingleObject(MainProcess, WIN32.INFINITE);
			Thread.Sleep(50);


			WIN32.CloseHandle(hReadStdin);
			WIN32.CloseHandle(hWriteStdout);

			return true;
		}

		~JudgeProcessor()
		{
			WIN32.CloseHandle(MainProcess);
			WIN32.CloseHandle(MainThread);
			WIN32.CloseHandle(ReadStream);
			WIN32.CloseHandle(job);
		}

		void Thread1()//监视异常和内存
		{
			UInt32 result;
			UInt32 key;
			IntPtr p = Marshal.AllocCoTaskMem(30);
			UInt32 R;
			while (WIN32.GetQueuedCompletionStatus(IoPort, out result, out key, p, WIN32.INFINITE)) {
				switch (result) {
					case WIN32.JOB_OBJECT_MSG_ACTIVE_PROCESS_ZERO:
						goto exit;
					case WIN32.JOB_OBJECT_MSG_PROCESS_MEMORY_LIMIT:
						lock (__Status) {
							_Status = JudeResult.MemoryLimitExceeded;
						}
						break;
					case WIN32.JOB_OBJECT_MSG_ABNORMAL_EXIT_PROCESS:
						WIN32.GetExitCodeProcess(MainProcess, out R);
						lock (__Status) {
							_Status = JudeResult.RuntimeError;
						}
						SetRuntimeError(R);
						break;
				}
			}
		exit:
			Marshal.FreeCoTaskMem(p);
		}

		void SetRuntimeError(uint R)
		{
			_Error = R;
		}

		void Thread2()//监视时间
		{
			UInt32 WaitRresult = WIN32.WaitForSingleObject(MainProcess, (UInt32)_TimeLimit);
			if (WaitRresult == WIN32.WAIT_TIMEOUT) {
				lock (__Status) {
					_Status = JudeResult.TimeLimitExceeded;
					WIN32.TerminateJobObject(job, 0);
				}
			}
		}

		void Thread3()//监视输出
		{
			uint readed = 0;
			var Received = new byte[1024];
			while (WIN32.ReadFile(ReadStream, Received, 1024, out readed, IntPtr.Zero)) {
				AddBytes(Received, readed);
			}
		}

		void AddBytes(byte[] Buffer, uint ReadedsBytes)
		{
			_OutputData.AddRange(Buffer.Take((int)ReadedsBytes));
			if (_OutputData.Count > OutputLimit) {
				lock (__Status) {
					if (_Status == JudeResult.OK || _Status == JudeResult.TimeLimitExceeded) {
						_Status = JudeResult.OutputLimitExceeded;
						WIN32.TerminateJobObject(job, 0);
						try {
							Thread.CurrentThread.Abort();
						} catch { }
					}
				}
			}
		}
	}
}
